/*
 *  Copyright 2008-2009 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#pragma once

#include <cusp/format.h>
#include <cusp/coo_matrix.h>

// SpMV
#include <cusp/detail/device/spmv/coo_flat.h>
#include <cusp/detail/device/spmv/csr_vector.h>
#include <cusp/detail/device/spmv/dia.h>
#include <cusp/detail/device/spmv/ell.h>
#include <cusp/detail/device/spmv/hyb.h>

#include <cusp/multi_matrix.h>
#include <cusp/detail/device/spmv/multi/multi_multiply.h>

// SpMM
#include <cusp/detail/device/spmm/coo.h>

#include <thrust/scatter.h>
#include <thrust/gather.h>

namespace cusp
{
namespace detail
{
namespace device
{

//////////////////////////////////
// Dense Matrix-Vector Multiply //
//////////////////////////////////
//// TODO implement this for both row and column-major ordering
//template <typename Matrix,
//          typename Vector1,
//          typename Vector2>
//void multiply(const Matrix&  A,
//              const Vector1& B,
//                    Vector2& C,
//              cusp::array2d_format,
//              cusp::array1d_format,
//              cusp::array1d_format)
//{
//}

///////////////////////////////////
// Sparse Matrix-Vector Multiply //
///////////////////////////////////
template <typename Matrix,
         typename Vector1,
         typename Vector2>
void multiply(const Matrix&  A,
              const Vector1& B,
              Vector2& C,
              cusp::coo_format,
              cusp::array1d_format,
              cusp::array1d_format)
{
#ifdef CUSP_USE_TEXTURE_MEMORY
    cusp::detail::device::spmv_coo_flat_tex(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#else
    cusp::detail::device::spmv_coo_flat(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#endif
}

template <typename Matrix,
         typename Vector1,
         typename Vector2>
void multiply(const Matrix&  A,
              const Vector1& B,
              Vector2& C,
              cusp::csr_format,
              cusp::array1d_format,
              cusp::array1d_format)
{
#ifdef CUSP_USE_TEXTURE_MEMORY
    cusp::detail::device::spmv_csr_vector_tex(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#else
    cusp::detail::device::spmv_csr_vector(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#endif
}

template <typename Matrix,
         typename Vector1,
         typename Vector2>
void multiply(const Matrix&  A,
              const Vector1& B,
              Vector2& C,
              cusp::dia_format,
              cusp::array1d_format,
              cusp::array1d_format)
{
#ifdef CUSP_USE_TEXTURE_MEMORY
    cusp::detail::device::spmv_dia_tex(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#else
    cusp::detail::device::spmv_dia(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#endif
}

template <typename Matrix,
         typename Vector1,
         typename Vector2>
void multiply(const Matrix&  A,
              const Vector1& B,
              Vector2& C,
              cusp::ell_format,
              cusp::array1d_format,
              cusp::array1d_format)
{
#ifdef CUSP_USE_TEXTURE_MEMORY
    cusp::detail::device::spmv_ell_tex(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#else
    cusp::detail::device::spmv_ell(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#endif
}

template <typename Matrix,
         typename Vector1,
         typename Vector2>
void multiply(const Matrix&  A,
              const Vector1& B,
              Vector2& C,
              cusp::permutation_format,
              cusp::array1d_format,
              cusp::array1d_format)
{
    thrust::gather( A.values.begin(), A.values.end(), B.begin(), C.begin() );
}


template <typename Matrix,
         typename Vector1,
         typename Vector2>
void multiply(const Matrix&  A,
              const Vector1& B,
              Vector2& C,
              cusp::hyb_format,
              cusp::array1d_format,
              cusp::array1d_format)
{
#ifdef CUSP_USE_TEXTURE_MEMORY
    cusp::detail::device::spmv_hyb_tex(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#else
    cusp::detail::device::spmv_hyb(A, thrust::raw_pointer_cast(&B[0]), thrust::raw_pointer_cast(&C[0]));
#endif
}

template <typename ValueType>
struct not_zero : public thrust::unary_function<ValueType,ValueType>
{
  __host__ __device__
  bool operator()(const ValueType x)
  {
    return x != ValueType(0);
  }
};

template <typename MultiMatrix,
         typename MultiVector1,
         typename MultiVector2>
void multiply(const MultiMatrix&  A,
              const MultiVector1& B,
              MultiVector2& C,
              cusp::multi_sparse_format,
              cusp::multi_array1d_format,
              cusp::multi_array1d_format)
{
    typedef typename MultiMatrix::index_type IndexType;
    typedef typename MultiMatrix::value_type ValueType;

    int num_slices = A.getNumSlices();

    if( thrust::detail::is_same<typename MultiMatrix::part_method, cusp::graph::colwise>::value )
    {
#pragma omp parallel for
        for ( int index = 0; index < num_slices; index++ )
        {
            int deviceNum = A.getDeviceNum(index);
            cudaSetDevice(deviceNum);
            cusp::multiply(A(index), B(index), C(index));
        }
        //cudaSetDevice(A.getDeviceNum(0));
        //cusp::blas::axpy(C(1),C(0),ValueType(1));

        /*cudaSetDevice(A.getDeviceNum(0));
	ValueType *y = (ValueType*) thrust::raw_pointer_cast(&A.temp_y[0][0]);
	thrust::copy( C(1).begin(), C(1).end(), thrust::device_pointer_cast(y) );
        cusp::blas::axpy(A.temp_y[0],C(0),ValueType(1));*/

	ValueType *y = (ValueType*) thrust::raw_pointer_cast(&A.temp_y[0][0]);

        cudaSetDevice(A.getDeviceNum(0));
	thrust::fill( thrust::device_pointer_cast(y), thrust::device_pointer_cast(y)+A.temp_y.size(), ValueType(0) );

        cudaSetDevice(A.getDeviceNum(1));
	thrust::scatter_if( C(1).begin(), C(1).end(), 
			    thrust::counting_iterator<IndexType>(0), 
			    thrust::make_transform_iterator(C(1).begin(), not_zero<ValueType>()),
			    thrust::device_pointer_cast(y) );
	cudaDeviceSynchronize();

        cudaSetDevice(A.getDeviceNum(0));
        cusp::blas::axpy(A.temp_y[0],C(0),ValueType(1));
    }
    else
    {
        if( A.num_rows != A.num_cols )
        {
#pragma omp parallel for
            for ( int index = 0; index < num_slices; index++ )
            {
                int deviceNum = A.getDeviceNum(index);
                cudaSetDevice(deviceNum);
                cusp::multiply(A(index), B(index), C(index));
            }
        }
        else
        {
            cusp::detail::device::multi_multiply( A, B, C,
                                                  typename MultiMatrix::multi_mult_method() );
        }
    }
}

////////////////////////////////////////
// Sparse Matrix-BlockVector Multiply //
////////////////////////////////////////
//// TODO implement this w/ repeated SpMVs and then specialize
//template <typename Matrix,
//          typename Vector1,
//          typename Vector2>
//void multiply(const Matrix&  A,
//              const Vector1& B,
//                    Vector2& C,
//              cusp::sparse_format,
//              cusp::array2d_format,
//              cusp::array2d_format)
//{
//}

////////////////////////////////////////
// Dense Matrix-Matrix Multiplication //
////////////////////////////////////////
// TODO implement
//template <typename Matrix1,
//          typename Matrix2,
//          typename Matrix3>
//void multiply(const Matrix1& A,
//              const Matrix2& B,
//                    Matrix3& C,
//              cusp::array2d_format,
//              cusp::array2d_format,
//              cusp::array2d_format)
//{
//}
template <typename Matrix1,
         typename Matrix2,
         typename Matrix3>
void multiply(const Matrix1& A,
              const Matrix2& B,
              Matrix3& C,
              cusp::permutation_format,
              cusp::coo_format,
              cusp::coo_format)
{
    cusp::detail::device::permute_rows(A,B,C);
}

template <typename Matrix1,
         typename Matrix2,
         typename Matrix3>
void multiply(const Matrix1& A,
              const Matrix2& B,
              Matrix3& C,
              cusp::coo_format,
              cusp::permutation_format,
              cusp::coo_format)
{
    cusp::detail::device::permute_columns(A,B,C);
}

/////////////////////////////////////////
// Sparse Matrix-Matrix Multiplication //
/////////////////////////////////////////
template <typename Matrix1,
         typename Matrix2,
         typename Matrix3>
void multiply(const Matrix1& A,
              const Matrix2& B,
              Matrix3& C,
              cusp::coo_format,
              cusp::coo_format,
              cusp::coo_format)
{
    cusp::detail::device::spmm_coo(A,B,C);
}

template <typename Matrix1,
         typename Matrix2,
         typename Matrix3>
void multiply(const Matrix1& A,
              const Matrix2& B,
              Matrix3& C,
              cusp::sparse_format,
              cusp::sparse_format,
              cusp::sparse_format)
{
    // other formats use COO * COO
    cusp::coo_matrix<typename Matrix1::index_type,typename Matrix1::value_type,cusp::device_memory> A_(A);
    cusp::coo_matrix<typename Matrix2::index_type,typename Matrix2::value_type,cusp::device_memory> B_(B);
    cusp::coo_matrix<typename Matrix3::index_type,typename Matrix3::value_type,cusp::device_memory> C_;

    cusp::detail::device::spmm_coo(A_,B_,C_);

    cusp::convert(C_, C);
}

/////////////////
// Entry Point //
/////////////////
template <typename Matrix,
         typename MatrixOrVector1,
         typename MatrixOrVector2>
void multiply(const Matrix&  A,
              const MatrixOrVector1& B,
              MatrixOrVector2& C)
{
    cusp::detail::device::multiply(A, B, C,
                                   typename Matrix::format(),
                                   typename MatrixOrVector1::format(),
                                   typename MatrixOrVector2::format());
}

} // end namespace device
} // end namespace detail
} // end namespace cusp

